﻿
#include <windows.h>
#include <shlobj.h>

#include <string>
#include <strsafe.h>

#include "Action/Script.h"
#include "resource.h"

#include "ContextMenu.h"
#include "ShellIcon.h"

HMODULE gModule = NULL;



// 类工厂
class MyContextMenuClassFactory : public IClassFactory {
public:
    MyContextMenuClassFactory() : m_refCount(1) {}

    HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppvObject) override {
        if (riid == IID_IUnknown || riid == IID_IClassFactory) {
            *ppvObject = static_cast<IClassFactory*>(this);
            AddRef();
            return S_OK;
        }
        *ppvObject = nullptr;
        return E_NOINTERFACE;
    }

    ULONG STDMETHODCALLTYPE AddRef() override {
        return InterlockedIncrement(&m_refCount);
    }

    ULONG STDMETHODCALLTYPE Release() override {
        ULONG refCount = InterlockedDecrement(&m_refCount);
        if (refCount == 0) {
            delete this;
        }
        return refCount;
    }

    HRESULT STDMETHODCALLTYPE CreateInstance(IUnknown* pUnkOuter, REFIID riid, void** ppvObject) override {
        if (pUnkOuter != nullptr) {
            return CLASS_E_NOAGGREGATION;
        }

        if (ppvObject == nullptr) {
            return E_POINTER;
        }

        IUnknown* pObject = nullptr;
        HRESULT hr = CLASS_E_CLASSNOTAVAILABLE;
        if (IsEqualCLSID(riid, IID_IContextMenu)) {
            pObject = reinterpret_cast<IUnknown*>(new MyContextMenu);
        }
        else if (IsEqualCLSID(riid, IID_IShellIconOverlayIdentifier)) {
            pObject = reinterpret_cast<IUnknown*>(new MyShellIcon);
        }
        else {
            return CLASS_E_CLASSNOTAVAILABLE;
        }

        if (pObject == nullptr) {
            return E_OUTOFMEMORY;
        }

        hr = pObject->QueryInterface(riid, ppvObject);
        pObject->Release();

        return hr;
    }

    HRESULT STDMETHODCALLTYPE LockServer(BOOL fLock) override {
        if (fLock) {
            InterlockedIncrement(&m_serverLocks);
        }
        else {
            InterlockedDecrement(&m_serverLocks);
        }
        return S_OK;
    }

public:
    LONG m_refCount;
    static LONG m_serverLocks;
};

LONG MyContextMenuClassFactory::m_serverLocks = 0;

// DLL 导出函数
HRESULT STDMETHODCALLTYPE DllGetClassObject(REFCLSID rclsid, REFIID riid, void** ppv) {
    static MyContextMenuClassFactory factory;
    return factory.QueryInterface(riid, ppv);
}

HRESULT STDMETHODCALLTYPE DllCanUnloadNow() {
    return (MyContextMenuClassFactory::m_serverLocks == 0) ? S_OK : S_FALSE;
}

HRESULT STDMETHODCALLTYPE DllRegisterServer() {
    MessageBox(NULL, TEXT("DllRegisterServer"), TEXT("DllRegisterServer"), MB_OK);
    HRSRC hRsrc = ::FindResource(gModule, MAKEINTRESOURCE(IDR_SCRIPT_INSTALL), TEXT("SCRIPT"));
    HGLOBAL hGlobal = ::LoadResource(gModule, hRsrc);
    DWORD dwSize = ::SizeofResource(gModule, hRsrc);

    CScript Script;
    Script.Init(gModule);
    Script.Parse((const char*)::LockResource(hGlobal));
    Script.Exec();

    if (hGlobal) UnlockResource(hGlobal);//其实可以不要这句
    // 注册 COM 组件的代码
    return S_OK;
}

HRESULT STDMETHODCALLTYPE DllUnregisterServer() {
    HRSRC hRsrc = ::FindResource(gModule, MAKEINTRESOURCE(IDR_SCRIPT_UNINSTALL), TEXT("SCRIPT"));
    HGLOBAL hGlobal = ::LoadResource(gModule, hRsrc);
    DWORD dwSize = ::SizeofResource(gModule, hRsrc);

    CScript Script;
    Script.Init(gModule);
    Script.Parse((const char*)::LockResource(hGlobal));
    Script.Exec();

    if (hGlobal) UnlockResource(hGlobal);//其实可以不要这句
    // 注销 COM 组件的代码
    return S_OK;
}

// DLL 入口点
BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved) {
    
    switch (ul_reason_for_call) {
    case DLL_PROCESS_ATTACH:
        gModule = hModule;
    case DLL_THREAD_ATTACH:
    case DLL_THREAD_DETACH:
    case DLL_PROCESS_DETACH:
        break;
    }
    return TRUE;
}
